import cv2
import numpy as np

NUM_H = 57
NUM_W = 88


def cv_show(name, img):
    """展示图片的函数"""
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def sort_contours(contours):
    """将轮廓按坐标位置从左向右排序"""
    # 轮廓里包含的都是该轮廓的点的坐标
    # 轮廓都是不一样的，需要用矩形框把每个轮廓框住，以左上角x坐标来排序
    bounds = [cv2.boundingRect(contour) for contour in contours]
    # lambda函数a[0]取每个元组(bound, contour)的bound，再[0]取bound第一个参数x (x, y, w, h)
    # 之后再zip很多个(bound, contour)得到一开始的(bounds, contours)格式
    bounds, contours = zip(*sorted(zip(bounds, contours), key=lambda a: a[0][0]))
    return contours


def get_digit2temp(template_img_path) -> dict:
    """获取数字映射到对应template模板图片的字典"""
    digit2temp = {}
    # 读入图片
    tmp = cv2.imread(template_img_path)
    cv_show("template", tmp)
    # 转换为灰度图
    # 这里不直接用灰度图读入是为了后面画轮廓的时候可以画彩色的线
    tmp_gray = cv2.cvtColor(tmp, cv2.COLOR_BGR2GRAY)
    cv_show("template_gray", tmp_gray)
    # 转换为二值图
    tmp_bin = cv2.threshold(tmp_gray, 10, 255, cv2.THRESH_BINARY_INV)[1]
    cv_show("template_bin_inv", tmp_bin)
    # 找每个数字的外轮廓
    contours = cv2.findContours(tmp_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
    # 在图上画出轮廓看看
    tmp_cp = tmp.copy()
    # 这里有个小坑点，之前怎么调color参数都是画的灰白色的线，后来才发现是画在了灰度图上面了，显示不了彩色
    cv2.drawContours(tmp_cp, contours, -1, (0, 255, 0), 3)
    cv_show("template_with_contours", tmp_cp)
    # 找出来的轮廓可能不是按数字的顺序排序的，需要手动排序一下
    contours_sort = sort_contours(contours)
    for n, contour in enumerate(contours_sort):
        x, y, w, h = cv2.boundingRect(contour)
        # cv2图片是 h w
        roi = tmp_bin[y:y + h, x:x + w]
        roi = cv2.resize(roi, (NUM_H, NUM_W))
        digit2temp[n] = roi
        cv_show(f"roi-{n}", roi)
    return digit2temp


def get_digits_bounds(card_img_path):
    """获取card图片中的连续数字(银行卡中数字一般四个一组)的矩形框"""
    # 先读入图片
    card = cv2.imread(card_img_path)
    cv_show("card", card)
    # 转灰度图
    card_gray = cv2.cvtColor(card, cv2.COLOR_BGR2GRAY)
    cv_show("card_gray", card_gray)
    # 定义卷积核
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 3))  # np.ones((9, 3), np.uint8)
    # 礼帽突出明亮区域
    card_tophat = cv2.morphologyEx(card_gray, cv2.MORPH_TOPHAT, kernel)
    cv_show("card_tophat", card_tophat)

    # # 梯度计算找出边缘
    # card_sobelx = cv2.Sobel(card_tophat, cv2.CV_64F, 1, 0, ksize=3)
    # card_sobelx = cv2.convertScaleAbs(card_sobelx)
    # card_sobely = cv2.Sobel(card_tophat, cv2.CV_64F, 0, 1, ksize=3)
    # card_sobely = cv2.convertScaleAbs(card_sobely)
    # card_sobel = cv2.addWeighted(card_sobelx, 0.5, card_sobely, 0.5, 0)
    # cv_show("card_sobel", card_sobel)
    # # 闭操作将数字连在一起(膨胀后腐蚀)
    # card_close = cv2.morphologyEx(card_sobel, cv2.MORPH_CLOSE, kernel)
    # card_close = cv2.morphologyEx(card_close, cv2.MORPH_CLOSE, kernel)
    # cv_show("card_close", card_close)

    # 发现膨胀两次更能保证将数字连上
    # 膨胀
    card_dilate = cv2.dilate(card_tophat, kernel, iterations=2)
    cv_show("card_dilate", card_dilate)
    # 闭操作将数字连在一起(膨胀后腐蚀)
    card_close = cv2.morphologyEx(card_dilate, cv2.MORPH_CLOSE, kernel)
    cv_show("card_close", card_close)
    # 转为二值图
    card_bin = cv2.threshold(card_close, 20, 255, cv2.THRESH_BINARY)[1]
    cv_show("card_bin", card_bin)
    # 找出轮廓
    contours = cv2.findContours(card_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
    card_cp = card.copy()
    # 画出轮廓
    cv2.drawContours(card_cp, contours, -1, (0, 255, 0), 2)
    cv_show("card_with_contours", card_cp)
    # 找出轮廓中需要的，用长宽的比例过滤
    bounds = []
    chose_contours = []
    for contour in contours:
        # 用矩形框框住
        bound = cv2.boundingRect(contour)
        x, y, w, h = bound
        pp = w / float(h)
        if 3.0 < pp < 4.5:
            # 还找了几个小框，在长度方面再限制
            if 100 < w < 200:
                bounds.append(bound)
                chose_contours.append(contour)
    # print(len(bounds))  # 4
    cv2.drawContours(card_cp, chose_contours, -1, (0, 0, 255), 2)
    cv_show("card_with_chose_contours", card_cp)
    # 矩形框从左到右排序
    bounds_sort = sorted(bounds, key=lambda a: a[0])
    return bounds_sort


def get_digit_result(group_bin, digit2temp, digit_bound):
    """获取单个数字的匹配结果"""
    scores = []
    for template in digit2temp.values():
        # 与模板中的每个数字进行匹配
        x, y, w, h = digit_bound
        digit_img = group_bin[y:y + h, x:x + w]
        digit_img = cv2.resize(digit_img, (NUM_H, NUM_W))
        # cv_show("digit", digit_img)
        result = cv2.matchTemplate(digit_img, template, cv2.TM_CCOEFF_NORMED)
        # 找到匹配该数字的最高分
        score = cv2.minMaxLoc(result)[1]
        scores.append(score)
    # 返回所有匹配数字分数的最高分的下标索引，即匹配的数字
    return str(np.argmax(scores))


def get_digit_results(img, digit2temp, digits_bounds):
    """获取所有数字的匹配结果"""
    digit_results = []
    for x, y, w, h in digits_bounds:
        # 拿到四个数字的方框
        group = img[y:y + h, x:x + w]
        cv_show("group", group)
        # 拆分每个数字
        # 转为二值图
        group_gray = cv2.cvtColor(group, cv2.COLOR_BGR2GRAY)
        group_bin = cv2.threshold(group_gray, 100, 255, cv2.THRESH_BINARY)[1]
        cv_show("group_bin", group_bin)
        # 找到轮廓
        contours = cv2.findContours(group_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        group_cp = group.copy()
        cv2.drawContours(group_cp, contours, -1, (0, 0, 255), 1)
        cv_show("group_with_contours", group_cp)
        # 用矩形框分割
        # 先要排序
        contours = sort_contours(contours)
        group_digit_bounds = [cv2.boundingRect(contour) for contour in contours]
        # print(len(digit_bounds))
        group_digit_results = [get_digit_result(group_bin, digit2temp, digit_bound)
                               for digit_bound in group_digit_bounds]
        # 画一下
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 0, 225), 2)
        cv2.putText(img, "".join(group_digit_results), (x, y-15), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2)
        digit_results += group_digit_results
    return digit_results


def identify(card_img_path, template_img_path):
    img = cv2.imread(card_img_path)
    digit2temp = get_digit2temp(template_img_path)
    digits_bounds = get_digits_bounds(card_img_path)
    digit_results = get_digit_results(img, digit2temp, digits_bounds)
    print("".join(digit_results))
    cv_show("result", img)


if __name__ == '__main__':
    identify(r"rsc/card01.png", "rsc/template.png")
